import cv2
import numpy as np
import math


def detect_yellow_ball():
    # 打开视频（可替换为摄像头0或其他视频路径）
    cap = cv2.VideoCapture("ping-pong.mp4")

    if not cap.isOpened():
        print("无法打开视频文件")
        return

    while True:
        ret, frame = cap.read()
        if not ret:
            print("视频读取完毕")
            break

        # 1. 转换为HSV颜色空间，提取黄色区域
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        # 调整黄色HSV范围（根据实际球体颜色微调）
        lower_yellow = np.array([20, 80, 80])  # 降低S和V的下限，适应不同亮度
        upper_yellow = np.array([35, 255, 255])  # 适当扩大H范围，避免漏检
        mask = cv2.inRange(hsv, lower_yellow, upper_yellow)

        # 2. 形态学处理：去除噪声，平滑轮廓
        kernel = np.ones((3, 3), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)  # 闭运算（填充小孔）
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)  # 开运算（去除小噪点）

        # 3. 查找轮廓
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # 4. 筛选圆形轮廓（黄色球体）
        for contour in contours:
            # 过滤过小的轮廓（排除噪声）
            area = cv2.contourArea(contour)
            if area < 50:  # 最小面积阈值（根据球体大小调整）
                continue

            # 计算周长（弧长）
            perimeter = cv2.arcLength(contour, closed=True)
            if perimeter == 0:  # 避免除以0
                continue

            # 计算圆形度
            circularity = 4 * math.pi * area / (perimeter ** 2)

            # 筛选圆形度较高的轮廓（例如 > 0.7，可根据实际调整）
            if circularity > 0.7:
                # 获取最小外接圆（更贴合球体）
                (x, y), radius = cv2.minEnclosingCircle(contour)
                center = (int(x), int(y))
                radius = int(radius)

                # 绘制圆和中心点
                cv2.circle(frame, center, radius, (0, 255, 255), 2)
                cv2.circle(frame, center, 2, (0, 0, 255), -1)  # 中心点
                cv2.putText(frame, "Yellow Ball", (center[0] - 50, center[1] - radius - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 2)

        # 显示结果
        cv2.imshow('Yellow Ball Detection', frame)
        cv2.imshow('Mask', mask)

        # 按'q'退出
        if cv2.waitKey(30) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    detect_yellow_ball()